import { auth, AuthResult, OAuthClientProvider, UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js'
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'
import { JSONRPCMessage, JSONRPCMessageSchema } from '@modelcontextprotocol/sdk/types.js'

/**
 * 自定义错误类，用于表示流式HTTP传输过程中的错误。
 */
export class StreamableHTTPError extends Error {
  constructor(
    public readonly code: number | undefined, // 错误码
    message: string | undefined, // 错误信息
    public readonly event: ErrorEvent // 相关事件对象
  ) {
    super(`Streamable HTTP error: ${message}`)
  }
}

/**
 * StreamableHTTPClientTransport 的配置项类型定义。
 */
export type StreamableHTTPClientTransportOptions = {
  /**
   * OAuth 客户端提供者，用于认证。
   *
   * 认证流程说明：
   * 1. 若指定了 authProvider，连接时会尝试使用已有的 access token。
   * 2. 若 token 过期，则尝试刷新 token。
   * 3. 若刷新失败或无 token，且服务端要求认证，则会调用 OAuthClientProvider.redirectToAuthorization，并抛出 UnauthorizedError。
   * 4. 用户授权后需调用 finishAuth 方法完成认证。
   */
  authProvider?: OAuthClientProvider

  /**
   * 自定义 HTTP 请求参数，可用于设置请求头等。
   */
  requestInit?: RequestInit
}

/**
 * 客户端流式HTTP传输实现，符合MCP Streamable HTTP协议规范。
 * 通过HTTP POST发送消息，通过HTTP GET（Server-Sent Events）接收消息。
 */
export class StreamableHTTPClientTransport implements Transport {
  private _activeStreams: Map<string, ReadableStreamDefaultReader<Uint8Array>> = new Map() // 活跃的SSE流
  private _abortController?: AbortController // 用于中止请求
  private _url: URL // 服务端URL
  private _requestInit?: RequestInit // 自定义请求参数
  private _authProvider?: OAuthClientProvider // OAuth认证提供者
  private _sessionId?: string // 会话ID
  private _lastEventId?: string // SSE最后事件ID

  // 事件回调
  onclose?: () => void
  onerror?: (error: Error) => void
  onmessage?: (message: JSONRPCMessage) => void

  /**
   * 构造函数，初始化传输对象。
   * @param url 服务端URL
   * @param opts 配置项
   */
  constructor(url: URL, opts?: StreamableHTTPClientTransportOptions) {
    this._url = url
    this._requestInit = opts?.requestInit
    this._authProvider = opts?.authProvider
  }

  /**
   * 内部方法：进行OAuth认证后启动连接。
   */
  private async _authThenStart(): Promise<void> {
    if (!this._authProvider) {
      throw new UnauthorizedError('No auth provider')
    }

    let result: AuthResult
    try {
      result = await auth(this._authProvider, { serverUrl: this._url })
    } catch (error) {
      this.onerror?.(error as Error)
      throw error
    }

    if (result !== 'AUTHORIZED') {
      throw new UnauthorizedError()
    }

    return await this._startOrAuth()
  }

  /**
   * 内部方法：生成通用请求头（包含认证信息和会话ID）。
   */
  private async _commonHeaders(): Promise<HeadersInit> {
    const headers: HeadersInit = {}
    if (this._authProvider) {
      const tokens = await this._authProvider.tokens()
      if (tokens) {
        headers['Authorization'] = `Bearer ${tokens.access_token}`
      }
    }

    if (this._sessionId) {
      headers['mcp-session-id'] = this._sessionId
    }

    return headers
  }

  /**
   * 内部方法：尝试建立SSE流或进行认证。
   */
  private async _startOrAuth(): Promise<void> {
    try {
      // 构建SSE请求头
      const commonHeaders = await this._commonHeaders()
      const headers = new Headers(commonHeaders)
      headers.set('Accept', 'text/event-stream')

      // 若有上次事件ID，则用于断点续传
      if (this._lastEventId) {
        headers.set('last-event-id', this._lastEventId)
      }

      // 尝试GET请求建立SSE流
      const response = await fetch(this._url, {
        method: 'GET',
        headers,
        signal: this._abortController?.signal
      })

      if (response.status === 405) {
        // 服务端不支持SSE，后续通过POST响应获取
        return
      }

      if (!response.ok) {
        if (response.status === 401 && this._authProvider) {
          // 需要认证
          return await this._authThenStart()
        }

        const error = new Error(`Failed to open SSE stream: ${response.status} ${response.statusText}`)
        this.onerror?.(error)
        throw error
      }

      // 成功建立SSE流，处理流数据
      const streamId = `initial-${Date.now()}`
      this._handleSseStream(response.body, streamId)
    } catch (error) {
      this.onerror?.(error as Error)
      throw error
    }
  }

  /**
   * 启动传输连接，建立SSE流。
   */
  async start() {
    if (this._activeStreams.size > 0) {
      throw new Error(
        'StreamableHTTPClientTransport already started! If using Client class, note that connect() calls start() automatically.'
      )
    }

    this._abortController = new AbortController()
    return await this._startOrAuth()
  }

  /**
   * 用户授权后调用此方法，完成OAuth认证流程。
   * @param authorizationCode 授权码
   */
  async finishAuth(authorizationCode: string): Promise<void> {
    if (!this._authProvider) {
      throw new UnauthorizedError('No auth provider')
    }

    const result = await auth(this._authProvider, { serverUrl: this._url, authorizationCode })
    if (result !== 'AUTHORIZED') {
      throw new UnauthorizedError('Failed to authorize')
    }
  }

  /**
   * 关闭所有活跃流，终止会话。
   */
  async close(): Promise<void> {
    // 关闭所有活跃流
    for (const reader of this._activeStreams.values()) {
      try {
        reader.cancel()
      } catch (error) {
        this.onerror?.(error as Error)
      }
    }
    this._activeStreams.clear()

    // 中止所有请求
    this._abortController?.abort()

    // 若有会话ID，发送DELETE请求通知服务端终止会话
    if (this._sessionId) {
      try {
        const commonHeaders = await this._commonHeaders()
        const response = await fetch(this._url, {
          method: 'DELETE',
          headers: commonHeaders,
          signal: this._abortController?.signal
        })

        if (!response.ok) {
          // 405为服务端不支持显式终止会话，不抛错
          if (response.status !== 405) {
            const text = await response.text().catch(() => null)
            throw new Error(`Error terminating session (HTTP ${response.status}): ${text}`)
          }
        }
      } catch (error) {
        // 即使终止失败也要触发onclose
        this.onerror?.(error as Error)
      }
    }

    this.onclose?.()
  }

  /**
   * 发送JSONRPC消息，支持批量发送。
   * @param message 单条或多条JSONRPC消息
   */
  async send(message: JSONRPCMessage | JSONRPCMessage[]): Promise<void> {
    try {
      const commonHeaders = await this._commonHeaders()
      const headers = new Headers({ ...commonHeaders, ...this._requestInit?.headers })
      headers.set('content-type', 'application/json')
      headers.set('accept', 'application/json, text/event-stream')

      const init = {
        ...this._requestInit,
        method: 'POST',
        headers,
        body: JSON.stringify(message),
        signal: this._abortController?.signal
      }

      const response = await fetch(this._url, init)

      // 处理会话ID
      const sessionId = response.headers.get('mcp-session-id')
      if (sessionId) {
        this._sessionId = sessionId
      }

      if (!response.ok) {
        if (response.status === 401 && this._authProvider) {
          // 认证过期，尝试重新认证
          const result = await auth(this._authProvider, { serverUrl: this._url })
          if (result !== 'AUTHORIZED') {
            throw new UnauthorizedError()
          }
          // 递归重试
          return this.send(message)
        }

        const text = await response.text().catch(() => null)
        throw new Error(`Error POSTing to endpoint (HTTP ${response.status}): ${text}`)
      }

      // 202表示已接受，无需处理body
      if (response.status === 202) {
        return
      }

      // 获取请求ID，用于追踪响应
      const messages = Array.isArray(message) ? message : [message]
      const requestIds = messages
        .filter((msg) => 'method' in msg && 'id' in msg)
        .map((msg) => ('id' in msg ? msg.id : undefined))
        .filter((id) => id !== undefined)

      const hasRequests = requestIds.length > 0
      const contentType = response.headers.get('content-type')

      if (hasRequests) {
        if (contentType?.includes('text/event-stream')) {
          // 流式响应，处理SSE流
          const streamId = `req-${requestIds.join('-')}-${Date.now()}`
          this._handleSseStream(response.body, streamId)
        } else if (contentType?.includes('application/json')) {
          // 非流式响应，直接处理JSON
          const data = await response.json()
          const responseMessages = Array.isArray(data)
            ? data.map((msg) => JSONRPCMessageSchema.parse(msg))
            : [JSONRPCMessageSchema.parse(data)]

          for (const msg of responseMessages) {
            this.onmessage?.(msg)
          }
        }
      }
    } catch (error) {
      this.onerror?.(error as Error)
      throw error
    }
  }

  /**
   * 内部方法：处理SSE流，解析事件并分发消息。
   * @param stream SSE流对象
   * @param streamId 流唯一标识
   */
  private _handleSseStream(stream: ReadableStream<Uint8Array> | null, streamId: string): void {
    if (!stream) {
      return
    }

    // 获取流读取器
    const reader = stream.getReader()
    this._activeStreams.set(streamId, reader)
    const decoder = new TextDecoder()
    let buffer = ''

    // 异步处理流数据
    const processStream = async () => {
      try {
        while (true) {
          const { done, value } = await reader.read()
          if (done) {
            // 流关闭
            this._activeStreams.delete(streamId)
            break
          }

          buffer += decoder.decode(value, { stream: true })

          // 按SSE协议分割事件
          const events = buffer.split('\n\n')
          buffer = events.pop() || ''

          for (const event of events) {
            const lines = event.split('\n')
            let id: string | undefined
            let eventType: string | undefined
            let data: string | undefined

            // 解析SSE事件格式
            for (const line of lines) {
              if (line.startsWith('id:')) {
                id = line.slice(3).trim()
              } else if (line.startsWith('event:')) {
                eventType = line.slice(6).trim()
              } else if (line.startsWith('data:')) {
                data = line.slice(5).trim()
              }
            }

            // 更新最后事件ID
            if (id) {
              this._lastEventId = id
            }

            // 处理消息事件
            if (data) {
              // 默认事件类型为'message'
              if (!eventType || eventType === 'message') {
                try {
                  const message = JSONRPCMessageSchema.parse(JSON.parse(data))
                  this.onmessage?.(message)
                } catch (error) {
                  this.onerror?.(error as Error)
                }
              }
            }
          }
        }
      } catch (error) {
        this._activeStreams.delete(streamId)
        this.onerror?.(error as Error)
      }
    }

    processStream()
  }
}
